#include <CFGBuilder/CFGBuild.h>
#include <fstream>
#include <iostream>
#include <algorithm>
#include <CFGBuilder/fileUtils.h>
#include <CFGBuilder/stringUtils.h>

using namespace std;
using namespace sca;

/**
 * initialize CFGBuilder
 * @param targetPath the path of target binary file
 */
CFGBuilder::CFGBuilder(const string &_targetPath, const string &_tempFileDir) : targetPath(_targetPath),
                                                                                tempFileDir(_tempFileDir) {
}

/**
 * determine whether the target binary file exist
 * @return if existing, return true,otherwise return false
 */
bool CFGBuilder::isTargetExist() const {
    return isFileExist(targetPath);
}

/**
 * load the instruction of the target binary file into unordered_map
 */
void CFGBuilder::loadInstruction() {
    // 生成临时的反汇编的文件，处理完毕后删除该文件
    generateDisassemblyFile(targetPath, tempFileDir);

    ifstream in(tempFileDir + "/bin.S");
    // a temporary value to record one line of the disassembly file
    string line;

    string beginAddr;
    bool beforeIsComment = false;

    // read instructions one by one
    while (getline(in, line)) {
        // spilt the line in order to get the instruction and its address
        trim(line);
        string delimiter("\t");
        // 地址的结束索引
        size_t index1 = getNthSpecChar(line, delimiter, 1);
        // 指令的起始索引
        size_t index2 = getNthSpecChar(line, delimiter, 2);

        // 如果当前行不是指令，而是注释信息，不进行处理
        if (index1 == string::npos) {
            beforeIsComment = true;
            continue;
        }

        string address = line.substr(0, index1 - 1);

        if (beforeIsComment) {
            beginAddr = address;
            beforeIsComment = false;
        }

        // 如果当前指令为空指令,空指令设置为nop
        if (index2 == string::npos) {
            // 添加空指令
            shared_ptr<Instruction> instructionPtr = make_shared<Instruction>("nop");
            ins.push_back(instructionPtr);
            inst2Add[instructionPtr] = address;
            add2Inst[address] = instructionPtr;
            continue;
        }

        // 去除指令中的注释
        string instruction = line.substr(index2 + 1);
        instruction = removeComments(instruction);

        // 使用容器存储每条指令
        shared_ptr<Instruction> instructionPtr = make_shared<Instruction>(instruction);
        ins.push_back(instructionPtr);
        inst2Add[instructionPtr] = address;
        add2Inst[address] = instructionPtr;

        // if current instruction is a jump kind instruction, add the jump destination into the entryAddress
        if (getInstructionKind(instruction) != InstructionKind::OTHER) {
            // 获取指令的类型
            auto instKind = getInstructionKind(instruction);
            if (instKind == InstructionKind::JUMP) {
                jInstruction[address] = instructionPtr;
            } else if (instKind == InstructionKind::RET) {
                retInstruction[address] = instructionPtr;
                beginAddr2Ret[beginAddr].push_back(instructionPtr);
            } else if (instKind == InstructionKind::CALL) {
                callInstruction[address] = instructionPtr;
            }
        }

    }
//     delete the temporary file
//    deleteDir("./tmp");
}

bool CFGBuilder::buildBasicBlocksVec() {
    if (ins.empty()) {
        return false;
    }
    shared_ptr<BasicBlock> bb = make_shared<BasicBlock>();
    string beginAddress;
    size_t insNum = ins.size();

    // traverse each instruction, and if encountering an instruction with an entry address, add a new basic block and
    // record the beginning address of the basic block
    for (size_t i = 0; i < insNum; i++) {
        if (beginAddress.empty()) {
            beginAddress = inst2Add[ins[i]];
        }
        if (entryAdds.find(inst2Add[ins[i]]) == entryAdds.end()) {
            bb->addInstructionIntoBB(ins[i]);
        } else {
            if (bb->getNumOfInstructions() != 0) {
                bbs.push_back(bb);
                addr2bb[beginAddress] = bb;
                bb2addr[bb] = beginAddress;
            }
            bb.reset();
            bb = make_shared<BasicBlock>();
            beginAddress.clear();
            bb->addInstructionIntoBB(ins[i]);
            beginAddress = inst2Add[ins[i]];
        }
    }

    if (bb->getNumOfInstructions() != 0) {
        bbs.push_back(bb);
        addr2bb[beginAddress] = bb;
        bb2addr[bb] = beginAddress;
    }
    return true;
}

void CFGBuilder::loadEntryAdds() {
    // 处理跳转指令
    for (auto iter = jInstruction.begin(); iter != jInstruction.end(); iter++) {
        // 跳转指令的下一条指令为另一个基本块的起始指令
        auto curInstIter = std::find(ins.begin(), ins.end(), iter->second);
        auto nextInstIter = ++curInstIter;
        // 如果当前指令有下一条指令，添加下一条指令的地址到entryAdds
        if (nextInstIter != ins.end()) {
            entryAdds.insert(inst2Add[*nextInstIter]);
        }

        // 跳转指令的目标指令是基本块的起始指令
        std::string targetAddr = getAddressFromJ(*iter->second);
        if (!targetAddr.empty())
            entryAdds.insert(targetAddr);
    }

    // 处理ret指令
    for (auto iter = retInstruction.begin(); iter != retInstruction.end(); iter++) {
        // ret指令的下一条指令为另一个基本块的起始指令
        auto curInstIter = std::find(ins.begin(), ins.end(), iter->second);
        auto nextInstIter = ++curInstIter;
        // 如果当前指令有下一条指令，添加下一条指令的地址到entryAdds
        if (nextInstIter != ins.end()) {
            entryAdds.insert(inst2Add[*nextInstIter]);
        }
    }

    // 处理call指令
    for (auto iter = callInstruction.begin(); iter != callInstruction.end(); iter++) {
        // call指令的下一条指令为另一个基本块的起始指令
        auto curInstIter = std::find(ins.begin(), ins.end(), iter->second);
        auto nextInstIter = ++curInstIter;
        // 如果当前指令有下一条指令，添加下一条指令的地址到entryAdds
        if (nextInstIter != ins.end()) {
            entryAdds.insert(inst2Add[*nextInstIter]);
        }

        // call指令的目标指令是基本块的起始指令
        std::string targetAddr = getAddressFromCall(*iter->second);
        if (!targetAddr.empty())
            entryAdds.insert(targetAddr);
    }
}

void CFGBuilder::createDependencies() {
    // 遍历每一个基本块
    for (auto iter = bbs.begin(); iter != bbs.end(); iter++) {
        // 获取当前基本块的最后一条指令的类型
        shared_ptr<Instruction> lastInst = (*iter)->getTheLastInstruction();
        // 如果当前基本块没有指令，跳过，不进行处理
        if (!lastInst) {
            continue;
        }
        InstructionKind lastInstKind = getInstructionKind(*lastInst);
        // 添加jmp指令依赖
        if (lastInstKind == InstructionKind::JUMP) {
            // 添加jmp指令的下一条指令的基本块
            auto curInstIter = std::find(ins.begin(), ins.end(), lastInst);
            auto nextInstIter = ++curInstIter;
            if (nextInstIter == ins.end()) {
                continue;
            }
            auto nextInstAddr = inst2Add[*nextInstIter];
            auto nextBB = addr2bb[nextInstAddr];
            (*iter)->addNextBasicBlock(nextBB);

            // 添加jmp跳转目标的基本块
            string targetAddr = getAddressFromJ(*lastInst);
            shared_ptr<BasicBlock> targetBB = addr2bb[targetAddr];
            // 如果存在跳转的目标基本块,直接添加
            if (targetBB)
                (*iter)->addNextBasicBlock(targetBB);
//            else{
//                // 如果目标基本块在动态库中，添加一个新的空基本块来填充空位
//                shared_ptr<BasicBlock> emptyBB = make_shared<BasicBlock>();
//                bbs.insert(bbs.end(),emptyBB);
//                addr2bb[targetAddr] = emptyBB;
//                bb2addr[emptyBB] = targetAddr;
//                (*iter)->addNextBasicBlock(emptyBB);
//            }
        } else if (lastInstKind == InstructionKind::CALL) {
            // 获取call指令的下一个指令所属的基本块
            auto curInstIter = std::find(ins.begin(), ins.end(), lastInst);
//            auto nextInstIter = ++curInstIter;
//            if (nextInstIter == ins.end()) {
//                continue;
//            }
//            // 添加call指令的下一条指令的基本块
//            auto nextInstAddr = inst2Add[*nextInstIter];
//            auto nextBB = addr2bb[nextInstAddr];
//            (*iter)->addNextBasicBlock(nextBB);

            // 获取call指令的目标地址
            string targetAddr = getAddressFromCall(*(*curInstIter));
            shared_ptr<BasicBlock> targetBB = addr2bb[targetAddr];
            // 如果存在跳转的目标基本块,直接添加
            if (targetBB)
                (*iter)->addNextBasicBlock(targetBB);
//            else{
//                // 如果目标基本块在动态库中，添加一个新的空基本块来填充空位
//                shared_ptr<BasicBlock> emptyBB = make_shared<BasicBlock>();
//                bbs.insert(bbs.end(),emptyBB);
//                addr2bb[targetAddr] = emptyBB;
//                bb2addr[emptyBB] = targetAddr;
//                (*iter)->addNextBasicBlock(emptyBB);
//            }
        } else if (lastInstKind == InstructionKind::RET) {
            if(retInst2CallInsts.find(lastInst) != retInst2CallInsts.end()){
                auto callInsts = retInst2CallInsts[lastInst];

                for(auto callInstIter = callInsts.begin() ; callInstIter != callInsts.end() ; callInstIter++ ){
                    auto curCallInstIter = std::find(ins.begin(), ins.end(), *callInstIter);
                    auto nextInstIter = ++curCallInstIter;
                    if (nextInstIter == ins.end()) {
                        continue;
                    }
                    // 添加call指令的下一条指令的基本块
                    auto nextInstAddr = inst2Add[*nextInstIter];
                    auto nextBB = addr2bb[nextInstAddr];

                    (*iter)->addNextBasicBlock(nextBB);
                }
            }

        } else {
            // 添加指令的下一条指令的基本块
            auto curInstIter = std::find(ins.begin(), ins.end(), lastInst);
            auto nextInstIter = ++curInstIter;
            if (nextInstIter == ins.end()) {
                continue;
            }
            auto nextInstAddr = inst2Add[*nextInstIter];
            auto nextBB = addr2bb[nextInstAddr];
            (*iter)->addNextBasicBlock(nextBB);
        }
    }
}

void CFGBuilder::buildCallAndRetRelationship() {
    for (auto iter = ins.begin(); iter != ins.end(); iter++) {
        if (getInstructionKind(**iter) == InstructionKind::CALL) {
            Address beginddr = getAddressFromCall(**iter);
            if (beginAddr2Ret.find(beginddr) != beginAddr2Ret.end()) {
                auto retInsts = beginAddr2Ret[beginddr];
                for (auto retInstIter = retInsts.begin(); retInstIter != retInsts.end(); ++retInstIter) {
                    retInst2CallInsts[*retInstIter].push_back(*iter);
                }
            }
        }
    }
}


